package main

import (
	"container/heap"
	"fmt"
)

type ListNode struct {
	Val  int
	Next *ListNode
}

type Heap []*ListNode

func (h *Heap) Len() int {
	return len(*h)
}

func (h *Heap) Less(i, j int) bool {
	return (*h)[i].Val < (*h)[j].Val
}

func (h *Heap) Swap(i, j int) {
	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
}

func (h *Heap) Push(x interface{}) {
	node := x.(*ListNode)
	*h = append(*h, node)
}

func (h *Heap) Pop() (v interface{}) {
	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
	return
}

// 优先对列实现
func main() {
	ln1 := &ListNode{
		Val: 1,
		Next: &ListNode{
			Val: 4,
			Next: &ListNode{
				Val:  5,
				Next: nil,
			},
		},
	}

	ln2 := &ListNode{
		Val: 1,
		Next: &ListNode{
			Val: 3,
			Next: &ListNode{
				Val:  4,
				Next: nil,
			},
		},
	}

	ln3 := &ListNode{
		Val: 1,
		Next: &ListNode{
			Val: 6,
			Next: &ListNode{
				Val:  8,
				Next: nil,
			},
		},
	}

	result := mergeKLists([]*ListNode{ln1, ln2, ln3})
	printListNode(result)
}

func printListNode(node *ListNode) {
	for node != nil {
		fmt.Printf("%d ->", node.Val)
		node = node.Next
	}
}

func mergeKLists(lists []*ListNode) *ListNode {
	h := &Heap{}
	heap.Init(h)

	for _, node := range lists {
		if node != nil {
			heap.Push(h, node)
		}
	}

	var ret, tmp *ListNode
	for {
		if h.Len() == 0 {
			break
		}
		minNode := heap.Pop(h).(*ListNode)

		node := &ListNode{
			Val: minNode.Val,
		}
		if ret == nil || tmp == nil {
			ret = node
			tmp = node
		} else {
			tmp.Next = node
			tmp = tmp.Next
		}
		if minNode.Next != nil {
			heap.Push(h, minNode.Next)
		}
	}

	return ret
}
